from mmdet_dataset import MMdetImageDataset
from competition_dataset import HPCompetitionImageDataset
import fiftyone as fo
from fiftyone import ViewField as F


pkl_file = 'prediction_final_6_13.pkl'
json_file = 'submit'
thres = 0.35
thres = [0.6, 0.35, 0.35, 0.35, 0.35, 0.7, 0.25, 0.25, 0.25, 0.25]
dataset = fo.Dataset.from_dir(name=None,
                                dataset_type=MMdetImageDataset,
                                dataset_dir=pkl_file,
                                shuffle=False,
                                threshold=thres,
                                tags=["predictions"])
print(dataset)
# vis_dataset = fo.Dataset("predictions", _create=True)
# vis_dataset.delete()
# vis_dataset = dataset.match(F('predict.confidence')> 0.25).clone()
# for sample in dataset[:2]:
#     print(sample)
    
dataset.export(json_file, dataset_type=HPCompetitionImageDataset)

# vis_dataset.name = "predictions"
# vis_dataset.save()